import cv2
from ultralytics import YOLO

# 1. 加载 YOLO 模型（使用 YOLOv5 或者 YOLOv8）
model = YOLO("yolov5s.pt")  # 你可以替换为适合的模型，支持 YOLOv8，YOLOv5 等

# 2. 读取图像
image_path = 'your_image.jpg'  # 替换为你的图像路径
image = cv2.imread(image_path)

# 获取原始图像的尺寸
original_height, original_width = image.shape[:2]

# 3. 将图像调整为 416x416
resized_image = cv2.resize(image, (416, 416))

# 4. 使用 YOLO 模型进行推理
results = model(resized_image)

# 获取模型标记结果（bounding boxes）
boxes = results[0].boxes  # 获取检测框
labels = results[0].names  # 获取标签名称

# 5. 在 YOLO 处理后的图像上绘制检测框
for box in boxes:
    x1, y1, x2, y2 = map(int, box.xyxy[0])  # 获取坐标
    confidence = box.conf[0].item()  # 获取置信度
    label = labels[int(box.cls[0].item())]  # 获取标签

    # 绘制框和标签
    cv2.rectangle(resized_image, (x1, y1), (x2, y2), (255, 0, 0), 2)
    cv2.putText(resized_image, f'{label} {confidence:.2f}', (x1, y1-10),
                cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 2)

# 6. 恢复图像至原始尺寸
restored_image = cv2.resize(resized_image, (original_width, original_height))

# 7. 显示处理后的图像
cv2.imshow("Processed Image", restored_image)
cv2.waitKey(0)
cv2.destroyAllWindows()

# 8. 保存处理后的图像（如果需要）
cv2.imwrite("processed_image.jpg", restored_image)
